package cc.unmi;

import cc.unmi.Topic;
import cc.unmi.serialization.AvroDeserializer;
import com.twitter.bijection.Injection;
import com.twitter.bijection.avro.GenericAvroCodecs;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.specific.SpecificRecordBase;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

public class BijectionConsumer {

    private KafkaConsumer<String, byte[]> consumer = new KafkaConsumer<>(getProperties());


    private static String schemaStr = "{\n" +
            "    \"type\": \"record\",\n" +
            "    \"name\": \"User\",\n" +
            "    \"fields\": [\n" +
            "        {\"name\": \"id\", \"type\": \"int\"},\n" +
            "        {\"name\": \"name\",  \"type\": \"string\"},\n" +
            "        {\"name\": \"age\", \"type\": \"int\"}\n" +
            "    ]\n" +
            "}";

    public List<GenericRecord> receive(Topic topic) {
        Schema.Parser parser = new Schema.Parser();
        Schema schema = parser.parse(schemaStr);
        Injection<GenericRecord, byte[]> recordInjection = GenericAvroCodecs.toBinary(schema);

//        TopicPartition partition = new TopicPartition(topic.topicName, 0);
        consumer.subscribe(Collections.singletonList(topic.topicName));
//        consumer.assign(Collections.singletonList(partition));
        ConsumerRecords<String, byte[]> records = consumer.poll(Duration.ofSeconds(3));

        return StreamSupport.stream(records.spliterator(), false)
                .map(ConsumerRecord::value).map(v -> recordInjection.invert(v).get())
                .collect(Collectors.toList());
    }

    private Properties getProperties() {
        Properties props = new Properties();
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, Topic.SERVERS_CONFIG);
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "bijection-cs-001");
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                StringDeserializer.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                ByteArrayDeserializer.class.getName());
        return props;
    }
}
